import os
import numpy as np
import matplotlib.pyplot as plt

"""
Relation between Y coordinate and Distance in the dataset
"""

# 读取 YOLO 格式的数据集
def read_yolo_dataset(folder_path):
    y_centers = []
    distances = []
    for filename in os.listdir(folder_path):
        if filename.endswith(".txt") and filename != "classes.txt":
            with open(os.path.join(folder_path, filename), "r") as file:
                lines = file.readlines()
                for line in lines:
                    class_id, x_center, y_center, width, height = map(
                        float, line.strip().split()
                    )
                    # 计算目标的中心 y 值和左上角点到右下角点的距离
                    y = y_center
                    # distance = np.sqrt(width**2 + height**2)  # 计算左上角点到右下角点的距离
                    distance = np.sqrt(width * height)
                    y_centers.append(y)
                    distances.append(distance)
    return np.array(y_centers), np.array(distances)


# 计算拟合效果的误差
def evaluate_fit(y, y_fit):
    mse = np.mean((y - y_fit) ** 2)  # 均方误差
    mae = np.mean(np.abs(y - y_fit))  # 平均绝对误差
    return mse, mae


# 从文件夹中读取 YOLO 格式的数据集
folder_path = "/home/hw/dataset/my_dataset/labels"  # 替换成你的数据集文件夹路径
y_centers, distances = read_yolo_dataset(folder_path)

# 进行多项式拟合
degree = 1  # 多项式的阶数
coefficients = np.polyfit(y_centers, distances, degree)

# 构建拟合的多项式函数
poly_function = np.poly1d(coefficients)
print(poly_function)
# 计算拟合效果的误差
distances_fit = poly_function(y_centers)
mse, mae = evaluate_fit(distances, distances_fit)
# 打印拟合效果的误差
print("均方误差 (MSE):", mse)
print("平均绝对误差 (MAE):", mae)
# 绘制散点图和拟合曲线
plt.scatter(y_centers, distances, s=5, label="Data")
plt.plot(
    np.sort(y_centers), poly_function(np.sort(y_centers)), "r-", label="Fitted curve"
)
plt.xlabel("Y coordinate")
plt.ylabel("Distance")
plt.title("Relation between Y coordinate and Distance")
plt.legend()
plt.show()
